#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@File    : wes_parse_snp.py
@Author  : Bing Liang
@Email   : believer19940901@gmail.com
@Date    : 2025/10/31 16:27
@Description : 将SNP过滤结果解析为结构化文件
"""

from argparse import ArgumentParser, Namespace
from pathlib import Path
import pandas as pd


# --------
# 辅助函数
# --------
def _snp_parser(sample: str, in_file: Path) -> pd.DataFrame:
    """
    解析单个样本的SNP过滤结果，并结构化为DataFrame。

    Args:
        sample: 样本名称
        in_file: 输入的SNP过滤结果文件路径

    Returns:
        pd.DataFrame: 解析后的数据框，包含SNP信息
    """
    # 读取文件内容，选取需要的列，忽略以#开头的注释行
    df = pd.read_csv(in_file, sep="\t", comment="#", header=None, usecols=[0, 1, 3, 4, 7, 8, 9])
    df.columns = ["Chr", "Pos", "Ref", "Alt", "ANN", "Info", "Stat"]

    # 输出字典，存储解析后的数据
    out_dict = {
        "Sample": [],  # 样本
        "AA_change": [],  # 突变模式
        "Chr": [],  # 染色体编号
        "Pos": [],  # 位点
        "Ref": [],  # 参考基因组
        "Alt": [],  # 变异基因组
        "Effect": [],  # 影响
        "Level": [],  # 等级
        "Gene_symbol": [],  # 基因符号
        "Mutation_type": [],  # 变异类型 (SNV 或 INDEL)
        "Depth": [],  # 总深度
        "Ref_depth": [],  # 参考位点深度
        "Alt_depth": [],  # 变异位点深度
        "Gene_type": []  # 纯合还是杂合
    }

    # 遍历每一行SNP数据
    for i, r in df.iterrows():
        c = r["Chr"]
        pos = r["Pos"]
        ref = r["Ref"]
        alt = r["Alt"]
        depth = r["Stat"].strip().split(":")[2]  # 获取总深度
        ref_depth = r["Stat"].strip().split(":")[1].strip().split(",")[0]  # 获取参考位点深度
        alt_depth = r["Stat"].strip().split(":")[1].strip().split(",")[1]  # 获取变异位点深度
        gene_flag = r["Stat"].strip().split(":")[0]  # 获取基因型标识

        # 判断基因型：1/1表示纯合(hom)，0/0表示纯合(hom)，其它表示杂合(het)
        gene_type = "hom" if gene_flag in ["1/1", "0/0"] else "het"

        # 解析ANN字段，获取变异的影响等信息
        ann = r['ANN'].strip().split(";")[-1].strip("ANN=")
        fields = ann.strip().split(",")

        # 遍历每个变异影响记录
        for field in fields:
            items = field.split("|")
            aa_change = f"{items[3]}:{items[6]}:{items[9]}"
            if items[10]:  # 如果存在蛋白质突变信息，加入
                aa_change += f":{items[10]}"
            effect = items[1]  # 影响
            level = items[2]  # 等级
            gene_symbol = items[4]  # 基因符号

            # 判断变异类型，如果参考和变异序列长度相等，表示为SNV，否则为INDEL
            mutation_type = "SNV" if len(ref) == len(alt) else "INDEL"

            # 填充字典
            out_dict["Sample"].append(sample)
            out_dict["AA_change"].append(aa_change)
            out_dict["Chr"].append(c)
            out_dict["Pos"].append(pos)
            out_dict["Ref"].append(ref)
            out_dict["Alt"].append(alt)
            out_dict["Effect"].append(effect)
            out_dict["Level"].append(level)
            out_dict["Gene_symbol"].append(gene_symbol)
            out_dict["Mutation_type"].append(mutation_type)
            out_dict["Depth"].append(depth)
            out_dict["Alt_depth"].append(alt_depth)
            out_dict["Ref_depth"].append(ref_depth)
            out_dict["Gene_type"].append(gene_type)

    return pd.DataFrame(out_dict)


# --------
# 主函数
# --------
def wes_parse_snp(args: Namespace):
    """
    解析多个样本的SNP过滤结果并输出到结构化的Excel文件。

    Args:
        args: 命令行参数，包含样本列表、输入文件路径和输出文件路径
    """
    # 解析输入参数，分割样本和文件路径
    samples = args.samples.strip().split(",")
    in_files = args.in_files.strip().split(",")

    # 输出文件路径
    out_file = Path(args.out_file).resolve()
    out_file.parent.mkdir(parents=True, exist_ok=True)  # 确保输出目录存在

    # 解析每个样本的SNP数据
    out_dfs = [_snp_parser(sample, in_file) for sample, in_file in zip(samples, in_files)]

    # 合并所有样本的数据框
    out_df = pd.concat(out_dfs, axis=0)

    # 将结果写入Excel文件
    out_df.to_excel(out_file, header=True, index=False, engine="openpyxl")


# --------
# 程序入口
# --------
if __name__ == '__main__':
    # 创建ArgumentParser实例，处理命令行参数
    parser = ArgumentParser(description="将SNP过滤结果解析为结构化文件")
    parser.add_argument(
        "--samples",
        type=str,
        required=True,
        help="所有样本的名称，逗号分隔"
    )
    parser.add_argument(
        "--in_files",
        type=str,
        required=True,
        help="所有样本的SNP过滤结果文件路径，逗号分隔"
    )
    parser.add_argument(
        "--out_file",
        type=str,
        required=True,
        help="输出文件的路径"
    )

    # 调用主函数进行处理
    wes_parse_snp(parser.parse_args())
